%%% This code generates the counterfactuals (robustness with respect to
%%% omega):

clear
load('counter.mat')

omega = -0.039; % Lower bound in Table 3, Allen and Donaldson (2022)
ubar_dyn = u_dyn.*(pop_dyn.^-omega);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%% cfa: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas:

pop_cfa = zeros(N,T); 
pop_cfa(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
prob_orig_given_dest_y_cf = prob_orig_given_dest_y_dyn(:,:,first_period);
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 0;

for t=first_period+1:first_period+T_aux-1
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfa(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_cf,intra_id,lambda_cf,alpha_dyn(:,first_period),squeeze(epsilon_dyn(:,:,t)),...
        pop_cfa(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfa = log(pop_cfa(:,2:end)) - log(pop_cfa(:,1:end-1));

%%% cfb: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas - intra model:

pop_cfb = zeros(N,T); 
pop_cfb(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 1;

for t=first_period+1:first_period+T_aux-1
    
    t
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfb(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_cf,intra_id,lambda_cf,alpha_intra_dyn(:,first_period),squeeze(epsilon_intra_dyn(:,:,t)),...
        pop_cfb(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfb = log(pop_cfb(:,2:end)) - log(pop_cfb(:,1:end-1));

%%% cfc: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas - model with no diffusion from migration:

pop_cfc = zeros(N,T); 
pop_cfc(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 0;
d_migr_const_aux_cfc = 0;

for t=first_period+1:first_period+T_aux-1
    
    t
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfc(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux_cfc,migration_exposure_cf,intra_id,lambda_cf,alpha_nomigr_dyn(:,first_period),squeeze(epsilon_nomigr_dyn(:,:,t)),...
        pop_cfc(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfc = log(pop_cfc(:,2:end)) - log(pop_cfc(:,1:end-1));

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Export data to run regressions in Stata:

for_stata = [repmat(cz_list,T_aux-1,1), reshape(repmat([1910:20:2010],N,1),N*(T_aux-1),1),...
    reshape(pop_growth_cfa(:,first_period:first_period+T_aux-2),N*(T_aux-1),1), ...
    reshape(pop_growth_cfb(:,first_period:first_period+T_aux-2),N*(T_aux-1),1), ...
    reshape(pop_growth_cfc(:,first_period:first_period+T_aux-2),N*(T_aux-1),1)];
    
csvwrite('results_model\cf_for_stata_omega_0p039.csv',for_stata)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear
load('counter.mat')

omega = -0.392; % Lower bound in Table 3, Allen and Donaldson (2022)
ubar_dyn = u_dyn.*(pop_dyn.^-omega);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%% cfa: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas:

pop_cfa = zeros(N,T); 
pop_cfa(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
prob_orig_given_dest_y_cf = prob_orig_given_dest_y_dyn(:,:,first_period);
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 0;

for t=first_period+1:first_period+T_aux-1
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfa(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_cf,intra_id,lambda_cf,alpha_dyn(:,first_period),squeeze(epsilon_dyn(:,:,t)),...
        pop_cfa(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfa = log(pop_cfa(:,2:end)) - log(pop_cfa(:,1:end-1));

%%% cfb: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas - intra model:

pop_cfb = zeros(N,T); 
pop_cfb(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 1;

for t=first_period+1:first_period+T_aux-1
    
    t
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfb(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_cf,intra_id,lambda_cf,alpha_intra_dyn(:,first_period),squeeze(epsilon_intra_dyn(:,:,t)),...
        pop_cfb(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfb = log(pop_cfb(:,2:end)) - log(pop_cfb(:,1:end-1));

%%% cfc: Estimate effect of the alpha shock by computing the model
%%% throughout without moving the alphas - model with no diffusion from migration:

pop_cfc = zeros(N,T); 
pop_cfc(:,1:first_period) = pop_dyn(:,1:first_period);
L_k_cf = L_k_dyn(:,first_period); L_y_cf = squeeze(L_y_dyn(:,:,first_period));
lambda_cf = squeeze(lambda_dyn(:,:,first_period));
migration_exposure_cf = migration_exposure_dyn(:,:,first_period);

intra_id = 0;
d_migr_const_aux_cfc = 0;

for t=first_period+1:first_period+T_aux-1
    
    t
    
    L_k_cf_aux = L_k_cf;
    
    [pop_cfc(:,t),lambda_cf,L_k_cf,L_y_cf,L_o_cf,pi_cf] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux_cfc,migration_exposure_cf,intra_id,lambda_cf,alpha_nomigr_dyn(:,first_period),squeeze(epsilon_nomigr_dyn(:,:,t)),...
        pop_cfc(:,t-1),L_k_cf,L_y_cf,ubar_dyn(:,t),mu,theta,zeta,omega,fertility_dyn(t));
    
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_cf(i,j) = sum(pi_cf(i,j,:))*L_k_cf_aux(i) / sum(L_y_cf(j,:));
        end
    end
    migration_exposure_cf = prob_orig_given_dest_y_cf.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_cf.*(1-eye(N)),1);
    
end

pop_growth_cfc = log(pop_cfc(:,2:end)) - log(pop_cfc(:,1:end-1));

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Export data to run regressions in Stata:

for_stata = [repmat(cz_list,T_aux-1,1), reshape(repmat([1910:20:2010],N,1),N*(T_aux-1),1),...
    reshape(pop_growth_cfa(:,first_period:first_period+T_aux-2),N*(T_aux-1),1), ...
    reshape(pop_growth_cfb(:,first_period:first_period+T_aux-2),N*(T_aux-1),1), ...
    reshape(pop_growth_cfc(:,first_period:first_period+T_aux-2),N*(T_aux-1),1)];
    
csvwrite('results_model\cf_for_stata_omega_0p392.csv',for_stata)




